package utils.datastructureZ.binaryTree;

import java.util.ArrayList;
import java.util.Iterator;

import static utils.datastructureZ.binaryTree.RBNode.NodeColor.BLACK;
import static utils.datastructureZ.binaryTree.RBNode.NodeColor.RED;

public class RedBlackTree<ElementType extends Comparable<ElementType>> implements Iterable<ElementType> {

    private final RBNode<ElementType> NIL;
    protected     RBNode<ElementType> root;
    protected     RBNode<ElementType> minNode;
    protected     RBNode<ElementType> maxNode;

    private abstract class itr implements Iterator<ElementType> {

        protected RBNode<ElementType> now;
        protected RBNode<ElementType> nil;

        public itr() {
            now = nil = null;
        }

        @Override
        public boolean equals(Object o) {
            if(!this.getClass().equals(o.getClass())) return false;
            return now == ((itr)o).now;
        }

    }

    public class MidItr extends itr {

        private final static boolean LEFT     = false;
        private final static boolean RIGHT    = true;
        private final static boolean TODOWN   = false;
        private final static boolean FROMDOWN = true;

        private boolean lrDirection;
        private boolean udDirection;

        public MidItr() {
            super();
            lrDirection = RIGHT;
            udDirection = TODOWN;
        }

        public MidItr(MidItr o) {
            now = o.now;
            nil = o.nil;
            lrDirection = o.lrDirection;
            udDirection = o.udDirection;
        }

        @Override
        public boolean hasNext() {
            return now != nil;
        }

        @Override
        public ElementType next() {

            ElementType ret = now.data;

            while(true) {
                if(udDirection == TODOWN) {
                    if(lrDirection == LEFT) {//这一步要去左下方
                        if(now.left != nil) {
                            now = now.left;
                        } else {
                            lrDirection = RIGHT;
                            return ret;
                        }
                    } else { //DOWN RIGHT
                        if(now.right != nil) {
                            now = now.right;
                            lrDirection = LEFT;
                        } else {
                            if(now == now.parent.left) lrDirection = LEFT;
                            udDirection = FROMDOWN;
                            now = now.parent;
                        }
                    }
                } else { //FROM DOWN
                    if(now == nil) return ret;
                    if(lrDirection == LEFT) { //这一步来自左下方
                        udDirection = TODOWN;
                        lrDirection = RIGHT;
                        return ret;
                    } else { //来自右下方
                        if(now == now.parent.left) lrDirection = LEFT;
                        now = now.parent;
                    }
                }
            }
        }

        /**
         * 按元素定位迭代器
         * @param root 树根节点
         * @param e    元素
         * @return     定位成功返回迭代器本身，定位失败返回null
         */
        MidItr locate(RBNode<ElementType> root, ElementType e) {
            if(root == NIL || e == null) return null;

            RBNode<ElementType> cur = root;

            while(cur != nil) {
                if(e.compareTo(cur.data) < 0) {
                    cur = cur.left;
                } else if(e.compareTo(cur.data) > 0) {
                    cur = cur.right;
                } else break;
            }

            if(cur != nil) {
                now = cur;
                lrDirection = RIGHT;
                if(cur.right != nil) udDirection = TODOWN;
                else                 udDirection = FROMDOWN;
                return this;
            }
            return null;
        }

        /**
         * 按照排名定位迭代器
         * @param root 树根节点
         * @param rank 排名
         * @return     定位成功返回迭代器本身，定位失败返回null
         */
        MidItr locateByRank(RBNode<ElementType> root, long rank) {
            if(root == NIL || rank < 1L) return null;

            RBNode<ElementType> cur = root; //遍历指针
            long curRank = cur.lNumber + 1;
            while(cur != NIL) {
                if(curRank < rank)      {
                    cur = cur.right;
                    curRank += cur.lNumber + 1;
                } else if(curRank > rank) {
                    cur = cur.left;
                    curRank -= cur.rNumber + 1;
                } else break;
            }

            now = cur;

            lrDirection = RIGHT;

            if(cur.right != nil) udDirection = TODOWN;
            else                 udDirection = FROMDOWN;
            return this;
        }

        public boolean before(MidItr o) {
            return now != nil && (o.now == nil || now.data.compareTo(o.now.data) < 0);
        }

        public boolean until(MidItr o) {
            return now != nil && (o.now == nil || now.data.compareTo(o.now.data) <= 0);
        }
    }

    protected MidItr mop, med;	//中序遍历迭代器的起点和终点
    private MidItr mid_begin() {return mop;}
    private MidItr mid_end()   {return med;}

    private void initMidIterator() {
        mop = new MidItr();
        med = new MidItr();

        mop.nil = NIL;
        med.now = NIL;

        mop.lrDirection = MidItr.RIGHT;
        mop.udDirection = MidItr.TODOWN;
        //重置迭代器
        resetMidIterator();
    }

    private void resetMidIterator() {
        mop.now = minNode;
    }

    @Override
    public Iterator<ElementType> iterator() {
        return new MidItr(mid_begin());
    }

    private void iteratorsInit()  { initMidIterator(); }
    private void iteratorsReset() { resetMidIterator(); }

    public RedBlackTree() {
        NIL = new RBNode<>(null, null, null, null, BLACK, 0L, 0L);
        NIL.left   = NIL;
        NIL.right  = NIL;
        NIL.parent = NIL;

        root = minNode = maxNode = NIL;
        iteratorsInit();
    }

    private void lRotate(RBNode<ElementType> nowAt)	{
        RBNode<ElementType> nowParent = nowAt.parent;
        RBNode<ElementType> nowright = nowAt.right;
                                            //01 y = x.right

        nowAt.rNumber = nowright.lNumber;
        nowright.lNumber += nowAt.lNumber + 1;


        nowAt.right = nowright.left;		//02 x.right = y.left

        if(nowright.left != NIL)			//03 if y.left != T.nil
            nowright.left.parent = nowAt;	//04	y.left.p = x

        nowright.parent = nowParent;		//05 y.p = x.p

        if(nowParent == NIL)				//06 if x.p == T.nil
            root = nowright;				//07 	T.root = y
        else if(nowAt == nowParent.left)	//08 if x == x.p.left
            nowParent.left = nowright;		//09 	x.p.left = y
        else nowParent.right = nowright;	//10 else x.p.right = y

        nowright.left = nowAt;				//11 y.left = x
        nowAt.parent = nowright;			//12 x.p = y
    }

    private void rRotate(RBNode<ElementType> nowAt)	{
        RBNode<ElementType> nowParent = nowAt.parent;
        RBNode<ElementType> nowleft = nowAt.left;
                                            //01 y = x.left

        nowAt.lNumber = nowleft.rNumber;
        nowleft.rNumber += nowAt.rNumber + 1;

        nowAt.left = nowleft.right;		//02 x.left = y.right

        if(nowleft.right != NIL)			//03 if y.right != T.nil
            nowleft.right.parent = nowAt;	//04	y.right.p = x

        nowleft.parent = nowParent;		//05 y.p = x.p

        if(nowParent == NIL)				//06 if x.p == T.nil
            root = nowleft;				//07 	T.root = y
        else if(nowAt == nowParent.right)	//08 if x == x.p.right
            nowParent.right = nowleft;	//09 	x.p.right = y
        else nowParent.left = nowleft;	//10 else x.p.left = y

        nowleft.right = nowAt;			//11 y.right = x
        nowAt.parent = nowleft;			//12 x.p = y
    }

    //插入平衡修复过程
    private void insertFix(RBNode<ElementType> nowAt) {
        RBNode<ElementType> nowParent = NIL;
        RBNode<ElementType> uncle     = NIL;
        while((nowParent = nowAt.parent).color == RED) {
            if(nowParent == nowParent.parent.left) {//当插入点的双亲节点是左节点
                uncle = nowParent.parent.right;

                if(uncle.color == RED) {//情况1，将其双亲节点与双亲的双亲节点颜色翻转（此时双亲的双亲节点肯定为黑色，否则不满足性质4）
                    nowParent.color = uncle.color = BLACK;
                    nowParent.parent.color = RED;
                    nowAt = nowParent.parent;
                } else {//情况2、3
                    if(nowAt == nowParent.right) {//情况2
                        nowAt = nowParent;
                        lRotate(nowAt);
                    }

                    //情况3
                    nowParent.color = BLACK;
                    nowParent.parent.color = RED;
                    rRotate(nowParent.parent);
                }
            } else {//当插入点的双亲节点是右节点
                uncle = nowParent.parent.left;

                if(uncle.color == RED) {//情况1，将其双亲节点与双亲的双亲节点颜色翻转（此时双亲的双亲节点肯定为黑色，否则不满足性质4）
                    nowParent.color = uncle.color = BLACK;
                    nowParent.parent.color = RED;
                    nowAt = nowParent.parent;
                } else {//情况2、3
                    if(nowAt == nowParent.left) {//情况2
                        nowAt = nowParent;
                        rRotate(nowAt);
                    }

                    //情况3
                    nowParent.color = BLACK;
                    nowParent.parent.color = RED;
                    lRotate(nowParent.parent);
                }
            }
        }
        root.color = BLACK;
    }

    //插入成功返回nullptr，插入元素已存在返回该元素的地址
    public ElementType op_insert(ElementType _data) {
        RBNode<ElementType> parentNode = NIL; //插入点的双亲
        RBNode<ElementType> nowAt = root;	  //遍历指针

        ArrayList<RBNode<ElementType>> updateL = new ArrayList<>();
        ArrayList<RBNode<ElementType>> updateR = new ArrayList<>();

        while(nowAt != NIL)	{
            parentNode = nowAt;
            if(nowAt.data.compareTo(_data) < 0)	{
                updateR.add(nowAt);
                nowAt = nowAt.right;
            } else if(nowAt.data.compareTo(_data) > 0) {
                updateL.add(nowAt);
                nowAt = nowAt.left;
            } else return nowAt.data;
        }

        RBNode<ElementType> newNode = new RBNode<>(_data, NIL, NIL, parentNode, RED, 0, 0);

        if(parentNode == NIL) root = newNode; //树空，插入节点作为新根
        else if(_data.compareTo(parentNode.data) < 0) parentNode.left  = newNode;
        else                                          parentNode.right = newNode;

        for(RBNode<ElementType> node: updateL) { ++node.lNumber; }
        for(RBNode<ElementType> node: updateR) { ++node.rNumber; }

        insertFix(newNode);		//修复平衡

        if(minNode != NIL) {
            while(minNode.left != NIL)
                minNode = minNode.left;
        } else minNode = root;

        if(maxNode != NIL) {
            while(maxNode.right != NIL)
                maxNode = maxNode.right;
        } else maxNode = root;

        iteratorsReset();
        return null;
    }

    public boolean insert(ElementType _data) {
        return op_insert(_data) == null;
    }
/*
    //插入成功返回nullptr，插入元素已存在返回该元素的地址
    public ElementType op_insert_debug(ElementType _data) {
        RBNode<ElementType> parentNode = NIL; //插入点的双亲
        RBNode<ElementType> nowAt = root;	  //遍历指针

        ArrayList<RBNode<ElementType>> updateL = new ArrayList<>();
        ArrayList<RBNode<ElementType>> updateR = new ArrayList<>();

        int i = 0;

        while(nowAt != NIL)	{
            parentNode = nowAt;
            if(nowAt.data.compareTo(_data) < 0)	{
                updateR.add(nowAt);
                nowAt = nowAt.right;
            } else if(nowAt.data.compareTo(_data) > 0) {
                updateL.add(nowAt);
                nowAt = nowAt.left;
            } else return nowAt.data;
        }

        RBNode<ElementType> newNode = new RBNode<>(_data, NIL, NIL, parentNode, RED, 0, 0);

        if(parentNode == NIL) root = newNode; //树空，插入节点作为新根
        else if(_data.compareTo(parentNode.data) < 0) parentNode.left  = newNode;
        else                                          parentNode.right = newNode;

        for(RBNode<ElementType> node: updateL) { ++node.lNumber; }
        for(RBNode<ElementType> node: updateR) { ++node.rNumber; }

        System.out.println("/////////////////////");
        testMid();
        System.out.println("/////////////////////");
        testPri();

        System.out.println("///////////////////////////////////////////");

        insertFix(newNode);		//修复平衡

        System.out.println("/////////////////////");
        testMid();
        System.out.println("/////////////////////");
        testPri();

        if(minNode != NIL) {
            while(minNode.left != NIL)
                minNode = minNode.left;
        } else minNode = root;

        if(maxNode != NIL) {
            while(maxNode.right != NIL)
                maxNode = maxNode.right;
        } else maxNode = root;

        return null;
    }

    public boolean insert_debug(ElementType _data) {
        return op_insert_debug(_data) == null;
    }
*/
    protected RBNode<ElementType> getMaxChild(RBNode<ElementType> nowAt) {
        if(nowAt == null) return null;

        while(nowAt.right != NIL)
            nowAt = nowAt.right;

        return nowAt;
    }

    protected RBNode<ElementType> getMinChild(RBNode<ElementType> nowAt) {
        if(nowAt == null) return null;

        while(nowAt.left != NIL)
            nowAt = nowAt.left;

        return nowAt;
    }

    protected RBNode<ElementType> getMinChildForRemove(
        RBNode<ElementType> nowAt,
        ArrayList<RBNode<ElementType>> updateL
	) {
        if(nowAt == null) return null;

        while(nowAt.left != NIL) {
            updateL.add(nowAt);
            nowAt = nowAt.left;
        }
        return nowAt;
    }

    public RBNode<ElementType> searchNode(ElementType _data) {
        RBNode<ElementType> nowAt = root;		//遍历指针
        while(nowAt != NIL)	{
            if(nowAt.data.compareTo(_data) < 0)		 nowAt = nowAt.right;
            else if(nowAt.data.compareTo(_data) > 0) nowAt = nowAt.left;
            else break;
        }

        return nowAt;
    }

    public ElementType search(ElementType _data) {
        RBNode<ElementType> targetNode = searchNode(_data);
        if(targetNode != NIL) return targetNode.data;
        return null;
    }

    public long getRank(ElementType _data) {

        if(root == NIL) return 0;

        RBNode<ElementType> nowAt = root;
        long rank = root.lNumber + 1;
        while(nowAt != NIL)	{
            if(nowAt.data.compareTo(_data) < 0) {
                nowAt = nowAt.right;
                rank += nowAt.lNumber + 1;
            } else if(nowAt.data.compareTo(_data) > 0) {
                nowAt = nowAt.left;
                rank -= nowAt.rNumber + 1;
            } else break;
        }

        return rank;
    }

    public long getMaxRank() {

        if(root == NIL) return 0;

        RBNode<ElementType> nowAt = root;
        long rank = root.lNumber + 1;
        while(nowAt.right != NIL) {
            nowAt = nowAt.right;
            rank += nowAt.lNumber + 1;
        }
        return rank;
    }

    private RBNode<ElementType> searchNodeByRank(long rank) {
        if(root == NIL || rank < 1L) return NIL;

        RBNode<ElementType> cur = root; //遍历指针
        long curRank = cur.lNumber + 1;
        while(cur != NIL) {
            if(curRank < rank) {
                cur = cur.right;
                curRank += cur.lNumber + 1;
            } else if(curRank > rank) {
                cur = cur.left;
                curRank -= cur.rNumber + 1;
            } else break;
        }

        return cur;
    }

    public ElementType searchByRank(long rank) {
        RBNode<ElementType> targetNode = searchNodeByRank(rank);
        if(targetNode != NIL) return targetNode.data;
        return null;
    }

    public MidItr locate(ElementType e) {
        MidItr midItr = new MidItr();
        midItr.nil = NIL;
        return midItr.locate(root, e);
    }

    public MidItr locateByRank(long rank) {
        MidItr midItr = new MidItr();
        midItr.nil = NIL;
        return midItr.locateByRank(root, rank);
    }

    /*
    public TreeSet<ElementType> getByRank(long low, long high) {
        searchNodeByRank()
    }
    */
    public RBNode<ElementType> searchNodeForRemove(
        ElementType _data,
        ArrayList<RBNode<ElementType>> updateL,
        ArrayList<RBNode<ElementType>> updateR
	) {
        RBNode<ElementType> nowAt = root;		//遍历指针
        //自顶向下寻找目标节点
        while(nowAt != NIL)	{
            if(nowAt.data.compareTo(_data) < 0)	{
                updateR.add(nowAt);
                nowAt = nowAt.right;
            } else if(nowAt.data.compareTo(_data) > 0) {
                updateL.add(nowAt);
                nowAt = nowAt.left;
            } else break;
        }

        return nowAt;
    }

    public ElementType min() {
        if(minNode == NIL) return null;
        return minNode.data;
    }

    public ElementType max() {
        if(maxNode == NIL) return null;
        return maxNode.data;
    }

    private void removeFix(RBNode<ElementType> nowAt) {
        RBNode<ElementType> brother = NIL;
        while(nowAt != root  &&  nowAt.color == BLACK) {
            if(nowAt == nowAt.parent.left) {
                brother = nowAt.parent.right;

                if(brother.color == RED) { //case 1
                    RBNode.NodeColor c = nowAt.parent.color;
                    nowAt.parent.color = brother.color;
                    brother.color      = c;
                    lRotate(nowAt.parent);
                    brother = nowAt.parent.right;
                }

                if( //case 2
                    brother.left.color == BLACK &&
                    brother.right.color == BLACK
                ) {
                    brother.color = RED;
                    nowAt = nowAt.parent;
                } else { //case 3
                    if(brother.right.color == BLACK) {
                        brother.color = RED;
                        brother.left.color = BLACK;
                        rRotate(brother);
                        brother = nowAt.parent.right;
                    } //case 4

                    brother.color = nowAt.parent.color;
                    nowAt.parent.color = BLACK;
                    brother.right.color = BLACK;
                    lRotate(nowAt.parent);

                    nowAt = root;
                }
            } else {
                brother = nowAt.parent.left;

                if(brother.color == RED) {	//case 1
                    RBNode.NodeColor c = nowAt.parent.color;
                    nowAt.parent.color = brother.color;
                    brother.color      = c;
                    rRotate(nowAt.parent);
                    brother = nowAt.parent.left;
                }

                if(	//case 2
                    brother.right.color == BLACK &&
                    brother.left.color == BLACK
                ) {
                    brother.color = RED;
                    nowAt = nowAt.parent;
                } else { //case 3
                    if(brother.left.color == BLACK) {
                        brother.color = RED;
                        brother.right.color = BLACK;
                        lRotate(brother);
                        brother = nowAt.parent.left;
                    }							//case 4
                    brother.color = nowAt.parent.color;
                    nowAt.parent.color = BLACK;
                    brother.left.color = BLACK;
                    rRotate(nowAt.parent);

                    nowAt = root;
                }
            }
        }

        nowAt.color = BLACK;
    }

    public RBNode<ElementType> op_remove(
        RBNode<ElementType> target,
        ArrayList<RBNode<ElementType>> updateL,
        ArrayList<RBNode<ElementType>> updateR
	) {
        if(target == NIL) return null;

        RBNode.NodeColor targetColor = target.color;
        RBNode<ElementType> removedPosition = NIL;
        RBNode<ElementType> tempParent = target.parent;

        if(target.right == NIL) {
            for(RBNode<ElementType> node: updateL) { --node.lNumber; }
            for(RBNode<ElementType> node: updateR) { --node.rNumber; }

            removedPosition = target.left;

            if(target == tempParent.left)
                tempParent.left = removedPosition;
            else tempParent.right = removedPosition;

            removedPosition.parent = tempParent;
            if(target == minNode) minNode = tempParent;
            if(target == maxNode) maxNode = tempParent;
            if(target == root) root = NIL;
        } else if(target.left == NIL) {
            //case_2:
            for(RBNode<ElementType> node: updateL) { --node.lNumber; }
            for(RBNode<ElementType> node: updateR) { --node.rNumber; }
            removedPosition = target.right;

            if(target == tempParent.left)
                tempParent.left = removedPosition;
            else tempParent.right = removedPosition;

            removedPosition.parent = tempParent;
            if(target == minNode) minNode = removedPosition;
            if(target == maxNode) maxNode = removedPosition;
            if(target == root) root = NIL;
        } else {
            updateR.add(target);
            removedPosition = getMinChildForRemove(target.right, updateL);

            tempParent = removedPosition.parent;

            ElementType d = target.data;
            target.data   = removedPosition.data;
            removedPosition.data = d;

            target = removedPosition;
            targetColor = target.color;
            tempParent = target.parent;
			//goto case_2;
            for(RBNode<ElementType> node: updateL) { --node.lNumber; }
            for(RBNode<ElementType> node: updateR) { --node.rNumber; }
            removedPosition = target.right;

            if(target == tempParent.left)
                tempParent.left = removedPosition;
            else tempParent.right = removedPosition;

            removedPosition.parent = tempParent;
            if(target == minNode) minNode = removedPosition;
            if(target == maxNode) maxNode = removedPosition;
            if(target == root) root = NIL;
        }

        if(targetColor == BLACK) removeFix(removedPosition);

        iteratorsReset();
        return target;
    }

    public boolean remove(ElementType _data) {

        ArrayList<RBNode<ElementType>> updateL = new ArrayList<>();
        ArrayList<RBNode<ElementType>> updateR = new ArrayList<>();

        RBNode<ElementType> target = searchNodeForRemove(_data, updateL, updateR);
        target = op_remove(target, updateL, updateR);

        return target != null;
    }

    public RBNode<ElementType> getNil() {
        return NIL;
    }

    private void testPriShow(RBNode<ElementType> nowAt) {
        if(nowAt != NIL) {
            System.out.println("Left:\t" + nowAt.lNumber + "\nData:\t" + nowAt.data.toString() + "\nColor:\t" + ((nowAt.color == RED)?"Red":"Black") + "\nRight:\t" + nowAt.rNumber + '\n');
            testPriShow(nowAt.left);
            testPriShow(nowAt.right);
        }
    }

    private void testMidShow(RBNode<ElementType> nowAt) {
        if(nowAt != NIL) {
            testMidShow(nowAt.left);
            System.out.println("Left:\t" + nowAt.lNumber + "\nData:\t" + nowAt.data.toString() + "\nColor:\t" + ((nowAt.color == RED)?"Red":"Black") + "\nRight:\t" + nowAt.rNumber + '\n');
            testMidShow(nowAt.right);
        }
    }

    public void testMid() {
        testMidShow(root);
    }

    public void testPri() {
        testPriShow(root);
    }

    public static void main(String[] args) {
        RedBlackTree<Integer> t = new RedBlackTree<>();

        //1 2 3 4 5 6 7 8 9
        //2 3 4 5 6 7 8 9 12

        t.insert(5); //

        t.insert(3); //
        t.insert(9); //
        t.insert(2); //
        t.insert(4); //
        t.insert(7); //
        t.insert(12);//
        t.insert(6); //
        t.insert(8); //

        //t.remove(5);
        t.testPriShow(t.root);

        System.out.println();

        System.out.println('\n' + "rank: " + t.getRank(5));
        System.out.println('\n' + "maxRank: " + t.getMaxRank());

        System.out.println();

        RedBlackTree<Integer>.MidItr low  = t.locateByRank(3L);
        RedBlackTree<Integer>.MidItr high = t.locateByRank(8L);

        while(low.until(high)) {
            System.out.print(" " + low.next());
        }

        System.out.println();
    }
}
